# -*- coding: utf-8 -*-

# @Author: Chen Hua
# @Date  : 2019/3/1
# @Desc  : 周开户数据

import sys
import datetime
import time
from db.MySqlConn import DBSession, JobConfig, engine as eg, InvestorAccount
import service.job_config_service as js
from util import DateUtil
import tushare as ts
from conf.settings import TUSHARE_TOKEN

##tushare sdk
ts.set_token(TUSHARE_TOKEN)
pro = ts.pro_api()
##链接数据库
session = DBSession()
job_name = 'stock_account_open'

"""
2008年1月到2015年5月29的数据调用这个接口
https://tushare.pro/document/2?doc_id=165
2015年5月29之后的数据调用这个接口
https://tushare.pro/document/2?doc_id=164
"""


def process_new_api(begin_dt, end_date):
    for _ in range(3):  # 重试三次
        new_ia = None
        try:
            df = pro.stk_account(start_date=DateUtil.date_to_yyyyMMhh_str(begin_dt),
                                 end_date=DateUtil.date_to_yyyyMMhh_str(end_date))
            if df.empty:
                print("没有数据返回，请检查是否出错NEW...")
            # 直接保存到数据库 investor_account 表
            for index, row in df.iterrows():
                #print(row)
                weekly_new = int(row['weekly_new'] * 10000)
                total = int(row['total'] * 10000)  # 单位万户转为户
                dt = end_date.date()
                new_ia = InvestorAccount(weekly_new=weekly_new, total=total, dt=dt)
        except Exception as e:
            print(str(e))
            time.sleep(2)
        else:
            return new_ia


def process_old_api(begin_dt, end_date):
    for _ in range(3):  # 重试三次
        new_ia = None
        try:
            df = pro.stk_account_old(start_date=DateUtil.date_to_yyyyMMhh_str(begin_dt),
                                     end_date=DateUtil.date_to_yyyyMMhh_str(end_date))
            if df.empty:
                print("没有数据返回，请检查是否出错OLD...")
            # 直接保存到数据库 investor_account 表
            for index, row in df.iterrows():
                print(row)
                weekly_new = int(row['new_sh'] + row['new_sz'])
                total = int((row['total_sh'] + row['total_sz']) * 10000)  # 单位万户转为户
                dt = end_date.date()
                new_ia = InvestorAccount(weekly_new=weekly_new, total=total, dt=dt)
        except Exception as e:
            print(str(e))
            time.sleep(2)
        else:
            return new_ia





def stock_account_open_info_sync():
    try:
        split_dt = datetime.datetime(2015, 5, 29)  ##周五
        now = datetime.datetime.now()
        cur_friday = now + datetime.timedelta(days=4 - now.weekday())

        begin_dt = datetime.datetime(2009, 1, 5)  ##默认初始从2009-01-05 周一日开始
        ##从数据库获取上次已获取的数据日期
        spider_config = session.query(JobConfig).filter(JobConfig.job_name == job_name).one()
        if spider_config and spider_config.job_mode == 'inc':
            begin_dt = spider_config.last_data_at
        # 做时间上的冗余
        end_date = begin_dt + datetime.timedelta(days=4)  # 周一加5天=周五

        while end_date < cur_friday.date():
            if end_date > split_dt.date():  ##日期大于2015-05-29日，用新的API
                ai = process_new_api(begin_dt, end_date)
            else:
                ai = process_old_api(begin_dt, end_date)
            print("成功爬取 ", str(begin_dt), "~", str(end_date), "的新开户信息...")
            if ai:
                session.merge(ai)
                session.commit()
            begin_dt = begin_dt + datetime.timedelta(days=7)  ##日期各加7天
            end_date = end_date + datetime.timedelta(days=7)
            time.sleep(15)  # 睡眠一秒，因为接口限制 100次/分钟
        js.update_job_config(job_name, end_date)
        print('Stock investor new open 信息已保存到表 investor_account ...')
    except Exception as e:
        print('获取Stock account new open 信息失败.')
        print(str(e))
        sys.exit(1)
    finally:
        session.close()


if __name__ == "__main__":
    stock_account_open_info_sync()
